import torchvision
import numpy as np
from PIL import Image
from copy import deepcopy
import torch

from dataset.noisify_cifar import noisify_cifar100_asymmetric
from dataset.imbalance_cifar import get_imb_data
from dataset.imagenet32 import Imagenet32


class CIFAR10(torchvision.datasets.CIFAR10):
    nb_classes = 10

    def __init__(self, root='~/data', train=True, transform=None,
                 r_ood=0.2, r_id=0.2, r_imb=0.1, seed=0, asym=False):
        print(f'using CIFAR-{self.nb_classes}...')
        super().__init__(root, train=train, transform=transform)
        if train is False:
            return
        if r_imb > 0.:
            self.data, self.targets = get_imb_data(self.data, self.targets,
                                                   self.nb_classes, r_imb, seed)
            print(f'Built imbalanced dataset, r_imb={r_imb}')
        np.random.seed(seed)

        if r_ood > 0.:
            ids_ood = [i for i in range(len(self.targets)) if np.random.random() < r_ood]
            imagenet32 = Imagenet32(root='~/data/imagenet32', train=True)
            img_ood = imagenet32.data[np.random.permutation(range(len(imagenet32)))[:len(ids_ood)]]
            self.ids_ood = ids_ood
            self.data[ids_ood] = img_ood
            print(f'Mixing in OOD noise, r_ood={r_ood}')

            if r_id > 0.:
                self.original_targets = deepcopy(self.targets)
                ids_not_ood = [i for i in range(len(self.targets)) if i not in ids_ood]
                ids_id = [i for i in ids_not_ood if np.random.random() < (r_id / (1 - r_ood))]
                if asym:
                    if self.nb_classes == 10:
                        transition = {0: 0, 2: 0, 4: 7, 7: 7, 1: 1, 9: 1, 3: 5, 5: 3, 6: 6, 8: 8}
                        for i, t in enumerate(self.targets):
                            if i in ids_id:
                                self.targets[i] = transition[t]
                    else:
                        self.targets = noisify_cifar100_asymmetric(self.targets, r_id, seed)
                    print(f'Mixing in ID asym noise, r_id={r_id}')
                else:
                    for i, t in enumerate(self.targets):
                        if i in ids_id:
                            self.targets[i] = int(np.random.random() * self.nb_classes)
                    print(f'Mixing in ID noise, r_id={r_id}')
                self.ids_id = ids_id

    def __getitem__(self, idx):
        image, target = self.data[idx], self.targets[idx]
        image = Image.fromarray(image)
        img = self.transform(image)
        target = torch.tensor(target).long()
        return {'index': idx, 'data': img, 'label': target}


class CIFAR100(CIFAR10):
    base_folder = "cifar-100-python"
    url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
    filename = "cifar-100-python.tar.gz"
    tgz_md5 = "eb9058c3a382ffc7106e4002c42a8d85"
    train_list = [
        ["train", "16019d7e3df5f24257cddd939b257f8d"],
    ]

    test_list = [
        ["test", "f0ef6b0ae62326f3e7ffdfab6717acfc"],
    ]
    meta = {
        "filename": "meta",
        "key": "fine_label_names",
        "md5": "7973b15100ade9c7d40fb424638fde48",
    }

    nb_classes = 100
